from fastapi import APIRouter
from fastapi.responses import StreamingResponse
from langchain.memory import ConversationBufferMemory
from langchain_huggingface import HuggingFaceEmbeddings
from pydantic import BaseModel
from langchain_openai import ChatOpenAI
from langchain_core.prompts import PromptTemplate,ChatPromptTemplate
from langchain.agents import create_react_agent,AgentExecutor,Tool
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain.schema import Document
import json
import asyncio
from Tools.weather import weather_func
from Tools.get_rag_db import get_rag_db_sync,get_rag_db_async
from Tools.rag_func import rag_func
from Tools.map_ import map
from Tools.os_ import os_f
import base64
from typing import Union
import os
import shutil

from db import models
tools=[
    Tool(
        name="weather check",
        func=weather_func.run,
        description='当用户询问某个城市的天气时，使用该工具，它可以检测指定城市天气情况,注意在输入时直接输入城市名就可以。'
    ),
    Tool(
        name="get_rag_db",
        func=get_rag_db_sync,
        coroutine=get_rag_db_async,
        description='获取所有可用的知识库信息。返回一个列表，包含每个知识库的描述、路径和对应的文本内容。这个工具只用于获取知识库参数，不能直接回答问题。获取到知识库信息后，如果用户问题与知识库相关，需要调用rag_func进行查询并且将对应的rag_path，chunks作为参数传入rag_func中。'
    ),
    Tool(
        name="rag_func",
        func=rag_func,
        description='该工具作用为通过知识库查询与用户提问有关的信息。参数以字典发送传入,该工具要求将上文3个参数依次传入。这是实际进行知识库查询的工具。'
    ),
    Tool(
        name="adr",
        func=map.get_address,
        description='输入某个地址，会返回一个字典，值分别为该地点的具体位置，与坐标'
    ),
    Tool(
        name='driver',
        func=map.driver,
        description='注意该工具的输入一定要是经adr处理后的坐标，输入两地的坐标，会返回两地之间驾车的交通路线，参数以字典发送传入,我要求你的参数输入格式为{"origin":（此处为起点坐标,要求输入字符串）,"destination":（此处为终点坐标，要求输入字符串）}'
    ),
    Tool(
        name='tr',
        func=map.transit,
        description='注意该工具的输入一定要是经adr处理后的坐标，输入两地的坐标与城市名，会返回两地之间公共交通的交通路线，参数以字典发送传入,我要求你的参数输入格式为{"origin":（此处为起点坐标，要求输入字符串）,"destination":（此处为终点坐标，要求输入字符串）,"city":(此处为城市名)}'
    ),
    Tool(
        name='os_',
        func=os_f.run,
        description='该工具可以将生成的信息保存到本地，输入数据格式为一个字典，其中格式为{"title":(文件标题),"content":(文件内容),"action":(文件操作,创建新文件并写入用w,在原有文件后添加用a),"end":(文件后缀如txt,png等)}'
    )
]


router = APIRouter()

class Chat_Message(BaseModel):
    query: str
    history: list
    temperature: float
    top_p: float
    model: str
    net_allow_statu: bool
    context: Union[str,None]


class Rag(BaseModel):
    rag_name: str
    rag_description: str
    rag_text: str

#调用大模型(网络)
async def get_net_model(api_key,url,model_name,tools,temperature,top_p,query,history,context):
    chat_model=ChatOpenAI(
        api_key=api_key,
        base_url=url,
        model=model_name,
        temperature=temperature,
        top_p=top_p,
        streaming=True
    )
    template = """
    请尽可能好的回答以下问题。如果需要可以适当使用一些工具。
    这是用户传过来的文本信息\n
    {context}\n
    这是你的聊天历史\n
    {history}\n
    你有以下工具可用：\n
    {tools}\n
    请使用以下格式: \n
    Question: 需要回答的问题。\n
    Thought: 总是要考虑应该做什么以及使用哪些工具。\n
    Action: 应采取的行动，应为[{tool_names}]中的一个。\n
    Action Input: 行动的输入。\n
    Observation: 行动的结果。\n
    ...（这个Thought/Action/Action Input/Observation 过程重复零次或多次）。\n
    Thought: 我现在知道最终答案了。\n
    Final Answer: 对原问题的最终答案。\n
    开始！ \n
    Question: {input}\n
    Thought: {agent_scratchpad}\n
    """
    prompt = PromptTemplate.from_template(template)
    agent = create_react_agent(chat_model, tools, prompt, stop_sequence=["\nObserv"])

    agent_executor = AgentExecutor.from_agent_and_tools(agent=agent, tools=tools, verbose=True)
    input_data={'input':query,'history':history,'context':context}
    async for chunk in agent_executor.astream(input_data):
        if "output" in chunk:
            yield f"data: {json.dumps({'content': chunk['output'], 'type': 'text'},ensure_ascii=False)}\n\n"
            await asyncio.sleep(0.01)
        elif "intermediate_steps" in chunk:
            steps = chunk["intermediate_steps"]
            for step in steps:
                if isinstance(step, tuple) and len(step) >= 2:
                    action, observation = step[0], step[1]
                    if hasattr(action, 'log'):
                        thought = action.log
                        yield f"data: {json.dumps({'content': f'思考: {thought}', 'type': 'thought'})}\n\n"
                        await asyncio.sleep(0.01)

    yield "data: [DONE]\n\n"


async def get_local_model(api_key, url, model_name, temperature, top_p, query,history):
    chat_model = ChatOpenAI(
        api_key=api_key,
        base_url=url,
        model=model_name,
        temperature=temperature,
        top_p=top_p,
        streaming=True
    )
    template = """
    你和用户的聊天记录:{history}\n
    请尽可能好的回答以下问题。
    问题: {input}
    请直接给出有帮助的回答：
    """
    prompt = PromptTemplate.from_template(template)
    chain = prompt | chat_model
    async for chunk in chain.astream({"input": query,'history': history}):
        if hasattr(chunk, 'content'):
            content = chunk.content
            yield f"data: {json.dumps({'content': content, 'type': 'text'}, ensure_ascii=False)}\n\n"
            await asyncio.sleep(0.01)

    # 发送结束信号
    yield "data: [DONE]\n\n"

#与AI会话
@router.post("/chat")
async def chat(message: Chat_Message):
    context=message.context
    query = message.query
    history = message.history
    temperature = message.temperature
    top_p = message.top_p
    from_model=message.model
    try:
        if from_model == "net_model":
            openai_api_key='sk-kmqzsejzspyeuhzgrrajsxuzjtlpdlmhorymtfickhniittw'
            base_url = "https://api.siliconflow.cn/v1"
            model = "Qwen/Qwen2.5-7B-Instruct"
            return StreamingResponse(
                get_net_model(openai_api_key, base_url, model, tools, temperature, top_p, query,history,context),
                media_type="text/plain; charset=utf-8"
            )
        else:
            openai_api_key = "EMPTY"
            base_url = "http://127.0.0.1:10222/v1"
            model = "Qwen2___5-0___5B-Instruct"
            return StreamingResponse(
                get_local_model(openai_api_key, base_url, model, temperature, top_p, query,history),
                media_type="text/plain; charset=utf-8"
            )
    except Exception as e:
        return str(e)

#添加知识库信息
@router.post("/add_rag")
async def rag(rag: Rag):
    rag_text = rag.rag_text
    rag_name = rag.rag_name
    rag_description = rag.rag_description

    exit=await models.Rag_db.filter(rag_name=rag_name)
    if exit:
        return {
            'code' : 404,
            'msg' : '向量库已存在'
        }

    faiss_dir = "./faiss_db"
    os.makedirs(faiss_dir, exist_ok=True)
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=10,
                                                   separators=['\n\n', '\n', ',', '.', '，', '。'])
    chunks = text_splitter.split_text(rag_text)
    documents = [Document(page_content=chunk) for chunk in chunks]
    chunk_len = len(chunks)
    embeddings = HuggingFaceEmbeddings(
        model_name="/home/lonely_bullet/work_place/models/AI-ModelScope/bge-large-zh-v1___5"
    )
    vs = FAISS.from_documents(documents, embeddings)
    save_path = f"{faiss_dir}/{rag_name}"
    vs.save_local(save_path)

    new_rag_db=await models.Rag_db.create(rag_name=rag_name,description=rag_description,rag_path=save_path,chunks=json.dumps(chunks))
    return {
        'code': 200,
        'msg' : '添加成功'
    }

#获取所有知识库的名字
@router.get("/get_all_db")
async def get_all_db():
    db=await models.Rag_db.all().values()
    data=[rag.get('rag_name') for rag in db]
    return{
        'code': 200,
        'data': data
    }

#获取指定数据库具体信息
@router.get("/get_rag_db")
async def get_rag_db(name : str):
    db=await models.Rag_db.filter(rag_name=name).first().values()
    chunks=db.get('chunks')
    chunks=json.loads(chunks)
    str='\n'.join(chunks)
    description=db.get('description')
    return {
        'code': 200,
        'data': {
            'name': name,
            'text': str,
            'description': description
        }
    }

#删除指定数据库
@router.delete("/drop_rag_db")
async def drop_rag_db(name : str):
    await models.Rag_db.filter(rag_name=name).delete()
    shutil.rmtree(f'./faiss_db/{name}')
    return {
        'code': 200,
    }


